"""
By Jiyuan Liu (liujiyuan13@163.com), Jun. 28, 2021.
All rights reserved.
"""

import numpy as np
import mat73
import sklearn.preprocessing as prep
from sklearn.metrics import *
import matplotlib.pyplot as plt

# def metric(y_true, y_pred):
#     # y_true[np.argwhere(y_true!=0)] = 1
#     # y_pred[np.argwhere(y_pred!=0)] = 1
#     acc = accuracy_score(y_true, y_pred)
#     f1_micro, f1_macro = f1_score(y_true, y_pred, average='micro'), f1_score(y_true, y_pred, average='macro')
#
#     return acc, f1_micro, f1_macro

def metric(y_true, y_pred, labels= [1,0], percent=True):
    """
    1 is positive, 0 is negative
         pred pred
          1    0
    gt 1  tp   fn
    gt 0  fp   tn
    :param y_true: numpy.array
    :param y_pred: numpy.array
    :param labels: [positive, negative]
    :param percent: if true, convert result into percent
    :return:
    fpr: for negative samples, the ratio of them will be classified as positive
    fnr: for positive samples, the ratio of them will be classified as negative
    error_rate: the ratio of mis-classification
    """
    tp, fn, fp, tn = confusion_matrix(y_true, y_pred, labels=labels).ravel()
    fpr, fnr, error_rate = fp/(fp+tn), fn/(fn+tp), (fp+fn)/(tp+fn+fp+tn)
    if percent:
        fpr, fnr, error_rate = fpr*100, fnr*100, error_rate*100
    return fpr, fnr, error_rate

def load_data(data_path):
    """
    load malware traffic dataset
    :param data_path: xxx.mat
    :return:
    """
    assert '.mat' in data_path

    # load mat file
    data = mat73.loadmat(data_path)
    X, Y = [prep.normalize(x[0], norm='l2').T for x in data['X']], data['Y']

    # remove feature "offered cipher suite" and "selected cipher suite"
    tmp = np.arange(786).tolist()
    del tmp[6:346]
    del tmp[396:736]
    X[1] = X[1][:, tmp]

    # set malware traffic with label 1, and normal with 0
    assert np.min(Y) == 0
    Y[np.argwhere(Y != 0)] = 1

    return X, Y



def gen_split(gt, tr_ratio=0.8):

    assert isinstance(gt, np.ndarray)

    labels = np.unique(gt)

    ind_tr, ind_te = [], []
    for i in range(len(labels)):
        ind = np.argwhere(gt==labels[i]).squeeze()
        np.random.shuffle(ind)
        num_tr = int(len(ind)*tr_ratio)
        ind_tr.extend(ind[:num_tr])
        ind_te.extend(ind[num_tr:])

    return ind_tr, ind_te

def twin_plot(x, y1, y2):

    fig, ax1 = plt.subplots()

    color = 'tab:red'
    ax1.set_xlabel('time (s)')
    ax1.set_ylabel('loss', color=color)
    ax1.plot(x, y1, color=color)
    ax1.tick_params(axis='y', labelcolor=color)

    ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis

    color = 'tab:blue'
    ax2.set_ylabel('error_rate', color=color)  # we already handled the x-label with ax1
    ax2.plot(x, y2, color=color)
    ax2.tick_params(axis='y', labelcolor=color)

    fig.tight_layout()  # otherwise the right y-label is slightly clipped
    plt.show()
